import os
import pandas as pd
import numpy as np
import pywt


def wavelet_transform_features(column, column_name=''):
    """对单列数据进行小波变换，并每wt_len个数据点进行一次特征提取。"""
    size = len(column)

    features = []
    # 绝对均值
    absolute_mean_value = np.sum(np.fabs(column)) / size
    # 峰值
    max = np.max(column)
    # 均方根值
    root_mean_score = np.sqrt(np.sum(np.square(column)) / size)
    # 方根幅值
    Root_amplitude = np.square(np.sum(np.sqrt(np.fabs(column))) / size)
    # 歪度值
    skewness = np.sum(np.power((np.fabs(column) - absolute_mean_value), 3)) / size
    # 峭度值
    Kurtosis_value = np.sum(np.power(column, 4)) / size
    # 波形因子
    shape_factor = root_mean_score / absolute_mean_value
    # 脉冲因子
    pulse_factor = max / absolute_mean_value
    # 歪度因子
    skewness_factor = skewness / np.power(root_mean_score, 3)
    # 峰值因子
    crest_factor = max / root_mean_score
    # 裕度因子
    clearance_factor = max / Root_amplitude
    # 峭度因子
    Kurtosis_factor = Kurtosis_value / np.power(root_mean_score, 4)

    wp = pywt.WaveletPacket(column, wavelet='db3', mode='symmetric', maxlevel=3)
    aaa = wp['aaa'].data
    aad = wp['aad'].data
    ada = wp['ada'].data
    add = wp['add'].data
    daa = wp['daa'].data
    dad = wp['dad'].data
    dda = wp['dda'].data
    ddd = wp['ddd'].data
    ret1 = np.linalg.norm(aaa, ord=None)
    ret2 = np.linalg.norm(aad, ord=None)
    ret3 = np.linalg.norm(ada, ord=None)
    ret4 = np.linalg.norm(add, ord=None)
    ret5 = np.linalg.norm(daa, ord=None)
    ret6 = np.linalg.norm(dad, ord=None)
    ret7 = np.linalg.norm(dda, ord=None)
    ret8 = np.linalg.norm(ddd, ord=None)

    # 频域
    data_fft = np.fft.fft(column)
    Y = np.abs(data_fft)
    freq = np.fft.fftfreq(size, 1 / 50000)
    ps = Y ** 2 / size
    # 重心频率
    FC = np.sum(freq * ps) / np.sum(ps)
    # 均方频率
    MSF = np.sum(ps * np.square(freq)) / np.sum(ps)
    # 均方根频率
    RMSF = np.sqrt(MSF)
    # 频率方差
    VF = np.sum(np.square(freq - FC) * ps) / np.sum(ps)

    features += [
        (f'{column_name}_ret1', ret1),  # 1
        (f'{column_name}_ret2', ret2),  # 2
        (f'{column_name}_ret3', ret3),  # 3
        (f'{column_name}_ret4', ret4),  # 4
        (f'{column_name}_ret5', ret5),  # 5
        (f'{column_name}_ret6', ret6),  # 6
        (f'{column_name}_ret7', ret7),  # 7
        (f'{column_name}_ret8', ret8),  # 8
        (f'{column_name}_absolute_mean_value', absolute_mean_value),  # 9
        (f'{column_name}_max', max),  # 10
        (f'{column_name}_root_mean_score', root_mean_score),  # 11
        (f'{column_name}_Root_amplitude', Root_amplitude),  # 12
        (f'{column_name}_skewness', skewness),  # 13
        (f'{column_name}_Kurtosis_value', Kurtosis_value),  # 14
        (f'{column_name}_shape_factor', shape_factor),  # 15
        (f'{column_name}_pulse_factor', pulse_factor),  # 16
        (f'{column_name}_skewness_factor', skewness_factor),  # 17
        (f'{column_name}_crest_factor', crest_factor),  # 18
        (f'{column_name}_clearance_factor', clearance_factor),  # 19
        (f'{column_name}_Kurtosis_factor', Kurtosis_factor),  # 20
        (f'{column_name}_FC', FC),  # 21
        (f'{column_name}_MSF', MSF),  # 22
        (f'{column_name}_RMSF', RMSF),  # 23
        (f'{column_name}_VF', VF),  # 24

    ]

    return features


def process_csv(file_path):
    """读取并处理单个CSV文件，返回重塑后的DataFrame"""
    df = pd.read_csv(file_path, names=['x切削力', 'y切削力', 'z切削力', 'x振动', 'y振动', 'z振动', '声发射'], header=None)

    # 对每列进行小波变换并提取特征
    features_dicts = []
    for column in df.columns:
        column_features = wavelet_transform_features(df[column], column_name=column)
        features_dicts.append(column_features)

    # 重组特征字典为DataFrame
    features_data = []
    for d in features_dicts:
        for feature_name, value in d:
            features_data.append({'Feature': feature_name, 'Value': value})

    # 将特征数据转换为DataFrame
    features_df = pd.DataFrame(features_data)

    # 创建一个空的DataFrame，用于存储重塑后的数据
    reshaped_df = pd.DataFrame(columns=features_df['Feature'].unique())

    # 将长格式数据转换为宽格式
    for index, row in features_df.iterrows():
        reshaped_df.loc[0, row['Feature']] = row['Value']  # 假设所有数据都属于一个组
    return reshaped_df


def process_all_csv(directory):
    """遍历指定目录下的所有CSV文件并合并结果"""
    all_dfs = []
    csv_files = [f for f in os.listdir(directory) if f.endswith('.csv')]
    total_files = len(csv_files)
    for i, filename in enumerate(csv_files):
        print(f'Processing file {i + 1}/{total_files}: {filename}')  # 打印进度
        file_path = os.path.join(directory, filename)
        df = process_csv(file_path)
        all_dfs.append(df)

    # 合并所有DataFrame
    combined_df = pd.concat(all_dfs, ignore_index=True)
    return combined_df


def merge_data(df, label):
    label['mean'] = label[['flute_1', 'flute_2', 'flute_3']].mean(axis=1)
    label['label'] = 0
    label.loc[label['mean'] > 120, 'label'] = 2
    label.loc[(label['mean'] >= 70) & (label['mean'] <= 120), 'label'] = 1
    df['label'] = label['label']
    return df


if __name__ == '__main__':
    import sys

    # 使用函数
    num = sys.argv[1]  # 从命令行参数获取数字
    folder_name = f'c{num}'  # 构造文件夹名称
    label_file = f'c{num}_wear.csv'  # 构造标签文件名
    combined_df = process_all_csv(folder_name)
    label = pd.read_csv(label_file)
    merge_data(df=combined_df, label=label).to_csv(f'data_c{num}.csv', index=None)
